from enum import IntEnum
from typing import *
from array import array
from operator import mul, add, neg
import tkinter as tk
import argparse
import logging
import os.path
import ast
import sys
import io

from ply.lex import lex
from ply.yacc import yacc

logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger()

class AstMeta(type):
    def __new__(cls, name, bases, dict_):
        if '__annotations__' in dict_:
            _fields = tuple(dict_['__annotations__'].keys())
        else:
            _fields = ()
        return super().__new__(cls, name, bases, dict(dict_) | {'_fields': _fields, '__match_args__': _fields})

    def __prepare__(self, *args, **kwargs):
        return OrderedDict()

class Ast(ast.AST, metaclass=AstMeta): ...

class Module(Ast):
    body = list[Ast]

class Stmt(Ast): ...
class Expr(Ast): ...
class Operator(Ast): ...

class Add(Operator): ...
class Sub(Operator): ...
class Mul(Operator): ...

class Id(Expr):
    id: str

class Array(Expr):
    value: list[int]

class MakeArray(Expr):
    length: int

class BinaryOp(Expr):
    op: Operator
    left: Expr
    right: Expr

class Index(Expr):
    array: Expr
    index: Expr

class Expression(Stmt):
    expr: Expr

class InplaceAdd(Stmt):
    id: str

class Print(Stmt):
    value: Expr

class SPrint(Stmt):
    value: Expr

class Exit(Stmt): ...

class Declare(Stmt):
    id: str
    value: Expr | None

class Assign(Stmt):
    id = str
    value: Expr

class SetItem(Stmt):
    array: Expr
    index: Expr
    value: Expr

class Cyberspace(Stmt): ...

reserved = [
    'print',
    'sprint',
    'u8',
    'cyberspace',
    'exit'
]
reserved = dict(zip(reserved, map(str.upper, reserved)))

tokens = [
    'LBRACKET',
    'RBRACKET',
    'OR',
    'SEMI',
    'EQ',
    'ADD',
    'SUB',
    'MUL',
    'DOUBLEADD',
    'INT',
    'ID'
]
tokens.extend(reserved.values())

t_ignore = ' \t\n'

t_LBRACKET = r'\['
t_RBRACKET = r'\]'
t_OR = r'\|'
t_SEMI = r';'
t_EQ = r'='
t_ADD = r'\+'
t_SUB = r'-'
t_MUL = r'\*'
t_DOUBLEADD = r'\+\+'

def t_COMMENT(t):
    r'//.*'

def t_error(c):
    raise SyntaxError('Bad character: {!r}'.format(c))

def t_INT(t):
    r'\d+'
    t.value = int(t.value)
    return t

def t_ID(t):
    r'\w+'
    t.type = reserved.get(t.value, 'ID')
    return t

lexer = lex()

def log(f, p):
    if p[0] is None:
        logger.warning('%s %s', f, list(p))

def p_module(p):
    '''
    module : stmt SEMI
    | module stmt SEMI
    '''
    match p[1:]:
        case Stmt() as stmt, ';':
            p[0] = Module(
                body=[stmt]
            )
        case Module() as mod, Stmt() as stmt, ';':
            mod.body.append(stmt)
            p[0] = mod
    log('module', p)

def p_factor(p):
    '''
    factor : array
    | ID
    | LBRACKET INT RBRACKET
    | factor LBRACKET expr RBRACKET
    array : INT
    | array OR INT
    '''
    match p[1:]:
        case int() as i, :
            p[0] = Array(
                value=[i]
            )
        case Array() as arr, '|', int() as i:
            arr.value.append(i)
            p[0] = arr
        case Array() as arr, :
            p[0] = arr
        case '[', int() as length, ']':
            p[0] = MakeArray(
                length=length
            )
        case str() as id, :
            p[0] = Id(
                id=id
            )
        case Expr() as value, '[', Expr() as index, ']':
            p[0] = Index(
                array=value,
                index=index
            )
    log('factor', p)

def p_expr(p):
    '''
    term : factor
    | term MUL factor
    expr : term
    | expr ADD term
    | expr SUB term
    '''
    match p[1:]:
        case Expr() as expr, :
            p[0] = expr
        case Expr() as left, str() as op, Expr() as right:
            p[0] = BinaryOp(
                left=left,
                op={
                    '+': Add(),
                    '-': Sub(),
                    '*': Mul()
                }[op],
                right=right
            )
    log('expr', p)

def p_stmt(p):
    '''
    stmt : U8 ID
    | U8 ID EQ expr
    | ID EQ expr
    | ID DOUBLEADD
    | PRINT expr
    | SPRINT expr
    | factor LBRACKET expr RBRACKET EQ expr
    | expr
    | CYBERSPACE
    | EXIT
    '''
    match p[1:]:
        case 'u8', str() as id:
            p[0] = Declare(
                id=id,
                value=None
            )
        case 'u8', str() as id, '=', Expr() as value:
            p[0] = Declare(
                id=id,
                value=value
            )
        case str() as id, '=', Expr() as value:
            p[0] = Assign(
                id=id,
                value=value
            )
        case str() as id, '++':
            p[0] = InplaceAdd(
                id=id
            )
        case 'print', Expr() as value:
            p[0] = Print(
                value=value
            )
        case 'sprint', Expr() as value:
            p[0] = SPrint(
                value=value
            )
        case Expr() as array, '[', Expr() as index, ']', '=', Expr() as value:
            p[0] = SetItem(
                array=array,
                index=index,
                value=value
            )
        case Expr() as expr, :
            p[0] = Expression(
                expr=expr
            )
        case 'cyberspace', :
            p[0] = Cyberspace()
        case 'exit', :
            p[0] = Exit()
    log('stmt', p)

parser = yacc()

class Instruction(IntEnum):
    PRINT = 0
    SPRINT = 1
    LOAD_VAR = 2
    LOAD_INT = 3
    BUILD_ARRAY = 4
    MAKE_ARRAY = 5
    GET_ITEM = 6
    STORE_ITEM = 7
    STORE_VAR = 8
    BINARY_OP = 9
    CYBERSPACE = 10
    EXIT = 11

    ADD = 0
    SUB = 1
    MUL = 2

class Assembler(ast.NodeVisitor):
    code: list[tuple[int, int]]
    symtable: list[str]

    def __init__(self):
        self.symtable = []

    def assembly(self, module: Module) -> tuple[list[tuple[int, int]], int]:
        self.code = []
        self.visit(module)
        return self.code, len(self.symtable)
    
    def lookup(self, id: str) -> int:
        try:
            return self.symtable.index(id)
        except ValueError:
            raise SyntaxError('Unbound identifier: {!r}'.format(id))
    
    def visit_Module(self, node: Module) -> None:
        for node in node.body:
            self.visit(node)

    def visit_Cyberspace(self, node: Cyberspace) -> None:
        self.code.append((Instruction.CYBERSPACE, 0))

    def visit_Exit(self, node: Exit) -> None:
        self.code.append((Instruction.EXIT, 0))

    def visit_Id(self, node: Id) -> None:
        self.code.append((Instruction.LOAD_VAR, self.lookup(node.id)))
        
    def visit_SetItem(self, node: SetItem) -> None:
        self.visit(node.array)
        self.visit(node.index)
        self.visit(node.value)
        self.code.append((Instruction.STORE_ITEM, 0))

    def visit_Assign(self, node: Assign) -> None:
        self.visit(node.value)
        self.code.append((Instruction.STORE_VAR, self.lookup(node.id)))

    def visit_Declare(self, node: Declare) -> None:
        if node.id in self.symtable:
            raise SyntaxError('Variable declaration duplicated: {!r}'.format(node.id))
        if node.value is None:
            self.visit(Array(
                value=[0]
            ))
        else:
            self.visit(node.value)
        self.code.append((Instruction.STORE_VAR, len(self.symtable)))
        self.symtable.append(node.id)

    def visit_Print(self, node: Print) -> None:
        self.visit(node.value)
        self.code.append((Instruction.PRINT, 0))

    def visit_SPrint(self, node: SPrint) -> None:
        self.visit(node.value)
        self.code.append((Instruction.SPRINT, 0))

    def visit_Expression(self, node: Expression) -> None:
        self.visit(node.expr)

    def visit_Index(self, node: Index) -> None:
        self.visit(node.array)
        self.visit(node.index)
        self.code.append((Instruction.GET_ITEM, 0))

    def visit_Array(self, node: Array) -> None:
        for elem in reversed(node.value):
            self.code.append((Instruction.LOAD_INT, elem))
        self.code.append((Instruction.BUILD_ARRAY, len(node.value)))

    def visit_InplaceAdd(self, node: InplaceAdd) -> None:
        self.visit(Assign(
            id=node.id,
            value=BinaryOp(
                op=Add(),
                left=Id(
                    id=node.id
                ),
                right=Array(
                    value=[1]
                ))
        ))

    def visit_MakeArray(self, node: MakeArray) -> None:
        self.code.append((Instruction.MAKE_ARRAY, node.length))

    def visit_BinaryOp(self, node: BinaryOp) -> None:
        self.visit(node.left)
        self.visit(node.right)
        self.code.append((Instruction.BINARY_OP, {
            Add: Instruction.ADD,
            Sub: Instruction.SUB,
            Mul: Instruction.MUL
        }[type(node.op)]))

def array_mul(left: array, right: array) -> array:
    return array('i', (sum(map(mul, left, right)), ))

def array_add(left: array, right: array) -> array:
    if len(left) == 1:
        return array('i', array('i', (i + left[0] for i in right)))
    if len(right) == 1:
        return array('i', array('i', (i + right[0] for i in left)))
    if len(left) != len(right):
        raise RuntimeError
    return array('i', map(add, left, right))

def array_sub(left: array, right: array) -> array:
    return array_add(left, array('i', map(neg, right)))

def array_setitem(arr: array, index: array, value: array) -> None:
    if len(value) != 1:
        raise RuntimeError
    value = value[0]
    if len(index) == 0:
        raise RuntimeError
    if index[0] == 0:
        if len(index) == 1:
            for i in range(len(arr)):
                arr[i] = value
            return
        raise RuntimeError
    for i in index:
        if i <= 0 or i > len(arr):
            continue
        arr[i - 1] = value

def array_getitem(arr: array, index: array):
    ret = array('i')
    for i in index:
        if i <= 0 or i > len(arr):
            continue
        ret.append(arr[i - 1])
    return ret

def run_helang(code: list[tuple[int, int]], scope: list[array], file=sys.stdout) -> None:
    stack: list[array] = []
    for inst, arg in code:
        match inst:
            case Instruction.PRINT:
                print(' | '.join(map(str, stack.pop())), file=file)
            case Instruction.SPRINT:
                print(bytes(stack.pop()).decode(), file=file)
            case Instruction.LOAD_VAR:
                stack.append(scope[arg])
            case Instruction.LOAD_INT:
                stack.append(arg)
            case Instruction.BUILD_ARRAY:
                stack.append(array('i', (stack.pop() for _ in range(arg))))
            case Instruction.MAKE_ARRAY:
                stack.append(array('i', (0 for _ in range(arg))))
            case Instruction.GET_ITEM:
                index = stack.pop()
                arr = stack.pop()
                stack.append(array_getitem(arr, index))
            case Instruction.STORE_ITEM:
                value = stack.pop()
                index = stack.pop()
                arr = stack.pop()
                array_setitem(arr, index, value)
            case Instruction.STORE_VAR:
                scope[arg] = stack.pop()
            case Instruction.BINARY_OP:
                right = stack.pop()
                left = stack.pop()
                stack.append((
                    array_add,
                    array_sub,
                    array_mul
                )[arg](left, right))
            case Instruction.CYBERSPACE:
                print('''// Getting your location...
// Your location is UNITED STATES.
// Congratulations! You are in the Cyber Spaces!''', file=file)
            case Instruction.EXIT:
                exit()
                
def code2bytes(code: list[tuple[int, int]]) -> bytes:
    return b''.join(int.to_bytes(inst) + int.to_bytes(arg) for inst, arg in code)

def bytes2code(bytes: bytes) -> list[tuple[int, int]]:
    res = []
    for i in range(0, len(bytes), 2):
        res.append((bytes[i], bytes[i + 1]))
    return res

def eval_helang(s: str, file=sys.stdout):
    code, vars_count = Assembler().assembly(parser.parse(s))
    run_helang(code, [None for _ in range(vars_count)], file)

def interactive(banner='Helang Interactive Console', prompt='> ', fout=sys.stdout):
    print(banner)
    assembler = Assembler()
    scope = []
    while True:
        try:
            s = input(prompt)
            if not s:
                continue
            if s.strip()[-1] != ';':
                s += ';'
            code, vars_count = assembler.assembly(parser.parse(s))
            while len(scope) < vars_count:
                scope.append(None)
            run_helang(code, scope, fout)
        except SystemExit:
            exit()
        except BaseException as ex:
            print(*ex.args, sep='\n')

def editor():
    def button_click():
        s = io.StringIO()
        eval_helang(input.get('1.0', 'end-1c'), s)
        output.delete('1.0', 'end')
        output.insert('end', s.getvalue())

    root = tk.Tk()
    root.title('Cyber Editor')
    root.geometry('1024x768')

    button = tk.Button(root, text="运行", command=button_click)
    button.pack()
    input = tk.Text(root, width=64, height=20)
    input.pack()
    output = tk.Text(root, width=64, height=20)
    output.insert('end', '输出')
    output.pack()

    root.mainloop()

argparser = argparse.ArgumentParser(description='Cyber Programming Language -- Helang')
argparser.add_argument('filename', action='store', nargs='?')
argparser.add_argument('--shell', action='store_true')
argparser.add_argument('--compile', action='store_true')
argparser.add_argument('--run', action='store_true')
argparser.add_argument('--debug', action='store_true')
argparser.add_argument('--editor', action='store_true')

if __name__ == '__main__':
    args = argparser.parse_args()
    if args.shell:
        interactive()
    elif args.editor:
        editor()
    elif args.filename:
        path = os.path.splitext(args.filename)[0] + '.hec'
    else:
        raise ValueError('Expected commandline argument "filename"')
    if args.run and os.path.exists(path):
        with open(path, 'rb') as fin:
            vars_count = fin.read(1)
            code = fin.read()
            run_helang(bytes2code(code), [None for _ in range(int.from_bytes(vars_count))])
    elif args.compile or args.run:
        with open(args.filename, 'r', encoding='utf-8') as fin, \
            open(path, 'wb') as fout:
            code, vars_count = Assembler().assembly(parser.parse(fin.read()))
            fout.write(vars_count.to_bytes())
            fout.write(code2bytes(code))
        if args.run:
            run_helang(code, [None for _ in range(vars_count)])
    elif args.debug:
        with open(args.filename, 'r', encoding='utf-8') as fin:
            eval_helang(fin.read())